#include <iostream>
#include <string.h>
#include <openssl/x509.h>
#include <openssl/rand.h>
#include <openssl/err.h>
#include <PopupLibTypes.hpp>
#include <PopupOSAL.hpp>
#include "PopupSSL.hpp"
#include "PopupUtilities.hpp"

using namespace std;
using namespace Popup;

static bool randInitialized = false;
#define RSA_PLAIN_TEXT_BLOC_SIZE 100

//=============================================================================
// SSL context
//=============================================================================

SslContext::SslContext(SslContextType ctxtType) :
  rsaLocal(new RSAHelper()), rsaRemote(0),
  deleteRsaLocal(true), blowfish(ctxtType)
{
}

SslContext::SslContext(SslContextType ctxtType,
                                 RSAHelper *p_localRsaHelper) :
  rsaLocal(p_localRsaHelper), rsaRemote(0),
  deleteRsaLocal(false), blowfish(ctxtType)
{
}

SslContext::~SslContext()
{
  if (deleteRsaLocal && (rsaLocal != 0)) {
    delete rsaLocal;
  }
  if (rsaRemote != 0) {
    delete rsaRemote;
  }
}

//! Returns the public key raw buffer
size_t SslContext::getLocalRSApublicKey(const unsigned char **p_key) const
{
  return rsaLocal->getRawPublicKey(p_key);
}

//! Returns the blowfish key raw buffer
const BlowfishKey & SslContext::getBlowfishKey() const {
  return blowfish.getKey();
}

void SslContext::setRemoteRSApublicKey(const unsigned char *p_key,
                                            size_t p_size)
{
  rsaRemote = new RSAHelper(p_key, p_size);
}

void SslContext::setRemoteBlowfishKey(const BlowfishKey & p_key) {
  blowfish.setKey(p_key);
}

bool SslContext::blowfishEncrypt(void *p_out,
                                      size_t p_size) const
{
  return blowfish.encrypt((unsigned char*) p_out,
                          (const unsigned char*) p_out,
                          p_size);
}

bool SslContext::blowfishEncrypt(void *p_out,
                                      const void *p_in,
                                      size_t p_size) const
{
  return blowfish.encrypt((unsigned char*) p_out,
                          (const unsigned char*) p_in,
                          p_size);
}

bool SslContext::blowfishDecrypt(void *p_out, size_t p_size) const
{
  return blowfish.decrypt((unsigned char*)p_out,
                          (const unsigned char*)p_out, p_size);
}

bool SslContext::blowfishDecrypt(void *p_out, const void *p_in,
                                      size_t p_size) const
{
  return blowfish.decrypt((unsigned char*)p_out,
                          (const unsigned char*)p_in,
                          p_size);
}

bool SslContext::rsaEncrypt(void *p_out, size_t *p_sizeOut,
                                 const void *p_in, size_t p_sizeIn) const
{
  // Encrypt with remote host public key
  return rsaRemote->encrypt((const unsigned char*) p_in, p_sizeIn,
                            (unsigned char *) p_out, p_sizeOut);
}

bool SslContext::rsaDecrypt(void *p_out, size_t *p_sizeOut,
                                 const void *p_in, size_t p_sizeIn) const
{
  // Decrypt with local private key
  return rsaLocal->decrypt((const unsigned char*) p_in, p_sizeIn,
                           (unsigned char*) p_out, p_sizeOut);
}

//=============================================================================
// RSA
//=============================================================================

RSAHelper::RSAHelper() :
  key(RSA_new()), num(BN_new()), mode(RSA_MASTER), rawpubkey(0),
      rawpubkeySize(0)
{
  if (randInitialized == false) {
    char salt[128];
    PopupUtils::randomString(&salt[0], 128);
    RAND_seed(salt, strlen(salt));
    randInitialized = true;
  }
  BN_set_word(num, RSA_F4);
  RSA_generate_key_ex(key, 1024, num, NULL);

  rawpubkeySize = i2d_RSAPublicKey(key, &rawpubkey);
}

RSAHelper::RSAHelper(const unsigned char *p_publicKey, size_t p_keySize) :
  key(d2i_RSAPublicKey(NULL, &p_publicKey, p_keySize)), num(0),
      mode(RSA_SLAVE), rawpubkey(0), rawpubkeySize(0)
{
}

RSAHelper::~RSAHelper()
{
  if (key != 0) {
    RSA_free(key);
  }
  if (num != 0) {
    BN_free(num);
  }
  if (rawpubkey != 0) {
    delete rawpubkey;
  }
}

size_t RSAHelper::getRawPublicKey(const unsigned char **p_key) const
{
  *p_key = rawpubkey;
  return rawpubkeySize;
}

bool RSAHelper::encrypt(const unsigned char *p_buffin, size_t p_sizeIn,
                        unsigned char *p_buffout, size_t *p_sizeOut) const
{
  size_t _blocsize = RSA_size(key);

  // High estimation of the size required in order to write the first
  // block into the buffer
  size_t _encryptedSize = 0;
  size_t _remainingSize = p_sizeIn;

  for (size_t offsetin = 0, offsetout = 0;
      offsetin < p_sizeIn;
      offsetin += RSA_PLAIN_TEXT_BLOC_SIZE, offsetout += _blocsize) {
    size_t _size = 0;
    if (_remainingSize < RSA_PLAIN_TEXT_BLOC_SIZE) {
      _size = _remainingSize;
    } else {
      _size = RSA_PLAIN_TEXT_BLOC_SIZE;
    }
    _remainingSize -= _size;
    _size = RSA_public_encrypt(_size, &p_buffin[offsetin],
                               &p_buffout[offsetout], key, RSA_PKCS1_PADDING);
    if (_size == (size_t) -1) {
      return false;
    } else {
      _encryptedSize += _size;
    }
  }
  *p_sizeOut = _encryptedSize;

  return (_encryptedSize > 0);
}

bool RSAHelper::decrypt(const unsigned char *p_buffin, size_t p_sizeIn,
                        unsigned char *p_buffout, size_t *p_sizeOut) const
{
  size_t _offsetout = 0;
  size_t _blocsize = RSA_size(key);

  for (size_t _offsetin = 0; _offsetin < p_sizeIn; _offsetin += _blocsize) {
    size_t _decsize = RSA_private_decrypt(RSA_size(key), &p_buffin[_offsetin],
                                     &p_buffout[_offsetout], key,
                                     RSA_PKCS1_PADDING);
    if (_decsize == (size_t) -1) {
      return false;
    } else {
      _offsetout += _decsize;
    }
  }
  *p_sizeOut = _offsetout;

  return (_offsetout > 0);
}

//=============================================================================
// Blowfish
//=============================================================================

#define NULL_IVEC { 0, 0, 0, 0, 0, 0, 0, 0 }

BlowfishHelper::BlowfishHelper(SslContextType ctxt)
{
  if (ctxt == CLIENT_CONTEXT) {
    PopupUtils::randomString((char*)&seed[0], sizeof(BlowfishKey));
    BF_set_key(&key, sizeof(BlowfishKey), seed);
  }
}

void BlowfishHelper::setKey(const BlowfishKey & p_key)
{
  memcpy(&seed[0], &p_key[0], sizeof(BlowfishKey));
  memset(&key, 0, sizeof(BF_KEY));
  BF_set_key(&key, sizeof(BlowfishKey), seed);
}

const BlowfishKey & BlowfishHelper::getKey() const
{
  return seed;
}

bool BlowfishHelper::encrypt(unsigned char *p_buffOut,
                             const unsigned char *p_buffIn,
                             size_t p_size) const
{
  unsigned char ivec[8] = NULL_IVEC;
  BF_cbc_encrypt(p_buffIn, p_buffOut, p_size, &key, ivec, BF_ENCRYPT);
  return true;
}

bool BlowfishHelper::decrypt(unsigned char *p_buffOut,
                             const unsigned char *p_buffIn,
                             size_t p_size) const
{
  unsigned char ivec[8] = NULL_IVEC;
  BF_cbc_encrypt(p_buffIn, p_buffOut, p_size, &key, ivec, BF_DECRYPT);
  return true;
}
